#include <ros/ros.h>
#include <behaviortree_cpp_v3/bt_factory.h>
#include <behaviortree_cpp_v3/loggers/bt_cout_logger.h>
#include <geometry_msgs/Twist.h>
#include <sensor_msgs/LaserScan.h>
#include <tf2_ros/transform_listener.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.h>

// 全局变量存储传感器数据
float obstacle_distance = 3.5;  // 默认无障碍物
geometry_msgs::PoseStamped target_pose;
geometry_msgs::PoseStamped robot_pose;

// 激光雷达回调函数 TODO
void laserCallback(const sensor_msgs::LaserScan::ConstPtr& msg) {
    obstacle_distance = *std::min_element(msg->ranges.begin(), msg->ranges.end());
}

// TF 监听机器人当前位置 TODO
void updateRobotPose(tf2_ros::Buffer& tf_buffer) {
    try {
        geometry_msgs::TransformStamped transform = tf_buffer.lookupTransform("map", "base_link", ros::Time(0));
        robot_pose.pose.position.x = transform.transform.translation.x;
        robot_pose.pose.position.y = transform.transform.translation.y;
    } catch (tf2::TransformException& ex) {
        ROS_WARN("%s", ex.what());
    }
}

// 行为树节点：检查障碍物
class CheckObstacle : public BT::ConditionNode {
public:
    CheckObstacle(const std::string& name) : BT::ConditionNode(name, {}) {}

    BT::NodeStatus tick() override {
        if (obstacle_distance < 0.5) {
            ROS_INFO("Obstacle detected! Distance: %.2f", obstacle_distance);
            return BT::NodeStatus::SUCCESS;  // 有障碍物
        } else {
            return BT::NodeStatus::FAILURE; // 无障碍物
        }
    }
};

// 行为树节点：检查是否到达目标 TODO
class CheckGoalReached : public BT::ConditionNode {
public:
    CheckGoalReached(const std::string& name) : BT::ConditionNode(name, {}) {}

    BT::NodeStatus tick() override {
        double dx = target_pose.pose.position.x - robot_pose.pose.position.x;
        double dy = target_pose.pose.position.y - robot_pose.pose.position.y;
        double distance = sqrt(dx * dx + dy * dy);

        if (distance < 0.2) {
            ROS_INFO("Goal reached!");
            return BT::NodeStatus::SUCCESS;  // 到达目标
        } else {
            return BT::NodeStatus::FAILURE;   // 未到达
        }
    }
};

// 行为树节点：发布前进命令
class MoveForward : public BT::ActionNodeBase {
public:
    MoveForward(const std::string& name, ros::Publisher& cmd_pub) 
        : BT::ActionNodeBase(name, {}), cmd_pub_(cmd_pub) {}

    BT::NodeStatus tick() override {
        geometry_msgs::Twist cmd;
        cmd.linear.x = 0.3;  // 前进速度
        cmd_pub_.publish(cmd);
        ROS_INFO("Moving forward");
        return BT::NodeStatus::RUNNING;
    }

    void halt() override {
        geometry_msgs::Twist cmd;
        cmd_pub_.publish(cmd);  // 停止
    }

private:
    ros::Publisher& cmd_pub_;
};

// 行为树节点：避障
class AvoidObstacle : public BT::ActionNodeBase {
public:
    AvoidObstacle(const std::string& name, ros::Publisher& cmd_pub) 
        : BT::ActionNodeBase(name, {}), cmd_pub_(cmd_pub) {}

    BT::NodeStatus tick() override {
        geometry_msgs::Twist cmd;
        cmd.angular.z = 0.5;  // 左转避障
        cmd_pub_.publish(cmd);
        ROS_INFO("Avoiding obstacle");
        return BT::NodeStatus::RUNNING;
    }

    void halt() override {
        geometry_msgs::Twist cmd;
        cmd_pub_.publish(cmd);  // 停止
    }

private:
    ros::Publisher& cmd_pub_;
};

// 行为树节点：停止
class Stop : public BT::ActionNodeBase {
public:
    Stop(const std::string& name, ros::Publisher& cmd_pub) 
        : BT::ActionNodeBase(name, {}), cmd_pub_(cmd_pub) {}

    BT::NodeStatus tick() override {
        geometry_msgs::Twist cmd;
        cmd_pub_.publish(cmd);
        ROS_INFO("Stopping");
        return BT::NodeStatus::SUCCESS;
    }

private:
    ros::Publisher& cmd_pub_;
};

int main(int argc, char** argv) {
    ros::init(argc, argv, "bt_decision_node");
    ros::NodeHandle nh;

    // 订阅激光雷达和目标点
    ros::Subscriber laser_sub = nh.subscribe("/tb3_0/scan", 1, laserCallback);
    ros::Publisher cmd_pub = nh.advertise<geometry_msgs::Twist>("/cmd_vel", 1);

    // TF 监听
    tf2_ros::Buffer tf_buffer;
    tf2_ros::TransformListener tf_listener(tf_buffer);

    // 设置目标点（示例）
    target_pose.pose.position.x = 5.0;
    target_pose.pose.position.y = 0.0;

    // 注册行为树节点
    BT::BehaviorTreeFactory factory;
    factory.registerNodeType<CheckObstacle>("CheckObstacle");
    factory.registerNodeType<CheckGoalReached>("CheckGoalReached");
    factory.registerNodeType<MoveForward>("MoveForward", cmd_pub);
    factory.registerNodeType<AvoidObstacle>("AvoidObstacle", cmd_pub);
    factory.registerNodeType<Stop>("Stop", cmd_pub);

    // 定义行为树
    std::string bt_xml = R"(
        <root main_tree_to_execute="MainTree">
            <BehaviorTree ID="MainTree">
                <Sequence name="MainSequence">
                    <CheckGoalReached name="CheckGoal"/>
                    <Fallback name="ObstacleAvoidance">
                        <CheckObstacle name="CheckObstacle"/>
                        <Sequence name="AvoidanceSequence">
                            <AvoidObstacle name="AvoidObstacle"/>
                            <MoveForward name="MoveForward"/>
                        </Sequence>
                    </Fallback>
                    <MoveForward name="MoveForward"/>
                </Sequence>
            </BehaviorTree>
        </root>
    )";

    // 创建行为树
    auto tree = factory.createTreeFromText(bt_xml);
    BT::StdCoutLogger logger(tree);  // 打印行为树日志

    // 主循环
    ros::Rate rate(10);
    while (ros::ok()) {
        updateRobotPose(tf_buffer);  // 更新机器人位置
        tree.tickRoot();             // 执行行为树
        ros::spinOnce();
        rate.sleep();
    }

    return 0;
}